import os
import re
import sys
import click
import subprocess
from pyfiglet import Figlet
from colorama import Fore, Style, init
from termcolor import colored
from six import print_
from mistralai import Mistral

# 初始化 Mistral 客户端
s = Mistral(
    api_key=os.getenv("MISTRAL_API_KEY", ""),
)

# 初始化对话历史
conversation_history = [
    {
        "content": """1. You are a helpful and friendly assistant with expertise in PowerShell scripting and "
        "command line.\n2. Assume the user is using the operating system `win` unless otherwise " "specified.\n3. Use "
        "the `code block` syntax in markdown to encapsulate any part in responses " "that is code, YAML, JSON, or XML, "
        "but not tables.\n4. When encapsulating command line code, " "use '```powershell' if it's a PowerShell "
        "command.\n5. When generating CLI commands, " "never break a command into multiple lines. Instead, always list "
        "all parameters and arguments " "of the command on the same line.\n6. Please keep the response concise but to "
        "the point. Do " "not overexplain.\n7. Only answer my needs, do not provide extra commands, Only one command "
        "can be given.\n8. When answering questions without executing commands, use the '```powershell echo' "
        "command to respond to the question. \n9. Add UTF-8 when reading the file,""",
        "role": "system",
    }
]


def get_response(user_message):
    # 添加用户消息到对话历史
    conversation_history.append({"content": user_message, "role": "user"})

    # 创建聊天补全请求
    res = s.chat.complete(model="pixtral-large-latest", messages=conversation_history)

    # 处理响应
    if res is not None:
        assistant_message = res.choices[0].message.content
        # 添加助手消息到对话历史
        conversation_history.append({"content": assistant_message, "role": "assistant"})
        return assistant_message
    else:
        return "No response from the model."


def execute_command(command):
    # 使用正则表达式去除 '```powershell' 和 '```'
    command = re.sub(r'```powershell', '', command).strip()
    command = re.sub(r'```', '', command).strip()

    if command.startswith("echo"):
        result = subprocess.run(["powershell", "-Command", command], capture_output=True, text=True)
        print(Fore.BLUE + result.stdout.strip() + Style.RESET_ALL)
    elif command.startswith("cd"):
        # 处理 cd 命令
        try:
            target_dir = command.split()[1]
            os.chdir(target_dir)
            print(Fore.YELLOW + f"Current directory: {os.getcwd()}" + Style.RESET_ALL)
        except Exception as e:
            print(Fore.RED + f"Error: {e}" + Style.RESET_ALL)
    else:
        # 使用 subprocess 调用 PowerShell 命令
        result = subprocess.run(["powershell", "-Command", command], capture_output=True, text=True)
        if result.returncode != 0:
            error_message = f"Error: {result.stderr}"
            print(Fore.RED + error_message + Style.RESET_ALL)
            # 将错误信息发送给 AI
            assistant_response = get_response(error_message)
            # 如果 AI 建议执行命令，则执行命令
            execute_command(assistant_response)
        else:
            output = result.stdout.strip()
            print(Fore.YELLOW + output + Style.RESET_ALL)
            # 如果命令是读取文件内容，将内容发送给 AI
            if "Get-Content" in command:
                assistant_response = get_response(output)
                # print(Fore.BLUE + "Assistant: " + Style.RESET_ALL + assistant_response)
                execute_command(assistant_response)


def clear_screen():
    # 清理屏幕
    os.system('cls' if os.name == 'nt' else 'clear')
    # 清理聊天历史
    clear_history()


def clear_history():
    # 清理聊天历史
    global conversation_history
    conversation_history = [
        {
            "content": """1. You are a helpful and friendly assistant with expertise in PowerShell scripting and "
        "command line.\n2. Assume the user is using the operating system `win` unless otherwise " "specified.\n3. Use "
        "the `code block` syntax in markdown to encapsulate any part in responses " "that is code, YAML, JSON, or XML, "
        "but not tables.\n4. When encapsulating command line code, " "use '```powershell' if it's a PowerShell "
        "command.\n5. When generating CLI commands, " "never break a command into multiple lines. Instead, always list "
        "all parameters and arguments " "of the command on the same line.\n6. Please keep the response concise but to "
        "the point. Do " "not overexplain.\n7. Only answer my needs, do not provide extra commands, Only one command "
        "can be given.\n8. When answering questions without executing commands, use the '```powershell echo' "
        "command to respond to the question. ,""",
            "role": "system",
        }
    ]
    print(Fore.BLUE + "Assistant: " + Style.RESET_ALL + "Chat history cleared.")


@click.command()
def terminal():
    init()
    f = Figlet(font='slant')
    print_(colored(f.renderText('PowerShell Terminal'), 'cyan'))

    while True:
        current_path = os.getcwd()
        user_input = input(f"{current_path}> " + Style.NORMAL)
        if user_input.lower() in ["exit", "quit"]:
            print(Fore.BLUE + "拜拜!")
            break
        elif user_input.lower() in ["clean", "清理", "-c"]:
            clear_screen()
            continue
        elif user_input.lower() in ["clear history", "清理历史", "-ch"]:
            clear_history()
            continue

        assistant_response = get_response(user_input)

        execute_command(assistant_response)


if __name__ == "__main__":
    terminal()